import tensorflow as tf

import os
os.environ['TF_CPP_MIN_LOG_LEVEL']='3'


if __name__ == '__main__':
    gpus = tf.config.list_physical_devices(device_type='GPU')
    tf.config.experimental.set_virtual_device_configuration(
        gpus[0],
        [
            tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048),
            tf.config.experimental.VirtualDeviceConfiguration(memory_limit=2048)
        ]
    )
    my_gpus=tf.config.experimental.list_logical_devices(device_type='GPU')
    print(my_gpus)
